package com.fw.leetcode.heap;

import com.fw.leetcode.LeetCode;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.PriorityQueue;

/**
 * 373. Find K Pairs with Smallest Sums
 *
 * You are given two integer arrays nums1 and nums2 sorted in non-decreasing order and an integer k.
 *
 * Define a pair (u, v) which consists of one element from the first array and one element from the second array.
 *
 * Return the k pairs (u1, v1), (u2, v2), ..., (uk, vk) with the smallest sums.
 *
 * Example 1:
 *  Input: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
 *  Output: [[1,2],[1,4],[1,6]]
 *  Explanation: The first 3 pairs are returned from the sequence: [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]
 *
 * Example 2:
 *  Input: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
 *  Output: [[1,1],[1,1]]
 *  Explanation: The first 2 pairs are returned from the sequence: [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]
 *
 * Constraints:
 *  1 <= nums1.length, nums2.length <= 10^5
 *  -10^9 <= nums1[i], nums2[i] <= 10^9
 *  nums1 and nums2 both are sorted in non-decreasing order.
 *  1 <= k <= 10^4
 *  k <= nums1.length * nums2.length
 */
public class Num_0373 implements LeetCode {
    private interface Solution {
        List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k);

        default void assertEquals(int[] nums1, int[] nums2, int k, int[][] expected) {
            List<List<Integer>> result = kSmallestPairs(nums1, nums2, k);
            for (int i = 0; i < result.size(); i++) {
                List<Integer> entry = result.get(i);
                if (!Arrays.equals(entry.stream().mapToInt(Integer::intValue).toArray(), expected[i])) {
                    kSmallestPairs(nums1, nums2, k);
                }
            }
        }
    }

    private static class MyHeapSolution implements Solution {

        @Override
        public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) { // 堆：时O(klogk) 空O(k)
            /*
             * 查找和最小的 K 对数字
             * - nums1 | nums2 已顺序排列
             * - 堆
             * [0,1,2,3,4] [0,1,2,3,4]
             *
             * 用堆存放索引，比较用数组索引处元素和
             * 1.用 (nums1[i],nums2[j]) 推出下两个可能组合 (nums1[i+1],nums2[j]) (nums1[i],nums2[j+1])
             * - 即：(i,j) -> (i+1,j) (i,j+1)
             * - 初始仅需入堆 (0,0)
             * 2.重复问题：
             * - (0,0) 出堆后，入堆 (0,1) (1,0)
             * - (0,1) 出堆后，入堆 (0,2) (1,1)
             * - (1,0) 出堆后，入堆 (1,1) (2,0) <- 此时 与上步 (1,1) 重复
             * - 解决：
             * - - 只对 j+1，i 不动，这样右边部分就都缺失
             * - - 再增加对右边的 (i,0) 的初始化入堆补偿即可
             *                                         (0,0)
             *                    (0,1)                                     (1,0)
             *          (0,2)               (1,1)                (1,1)               (2,0)
             *     (0,3)     (1,2)      (1,2)    (2,1)       (1,2)    (2,1)       (2,1)    (3,0)
             * (0,4)(1,3) (1,3)(2,2) (1,3)(2,2)(2,2)(3,1) (1,3)(2,2)(2,2)(3,1) (2,2)(3,1)(3,1)(4,0)
             */
            List<List<Integer>> result = new ArrayList<>(k);
            // 小顶堆：排序大小比较  sum(nums1[p1[0]], nums2[p1[1]]) - sum(nums1[p2[0]], nums2[p2[1]])
            PriorityQueue<int[]> heap = new PriorityQueue<>(k, (p1, p2)->
                    nums1[p1[0]] + nums2[p1[1]] - nums1[p2[0]] - nums2[p2[1]]); // 小心顺序
            heap.add(new int[]{0,0}); // 初始起点：最小值 => (nums1[0],nums2[0]) =(sum)=> nums1[0]+nums2[0]
            for (int i = 1; i < nums1.length; i++) { // [去重优化点] 避免重复后的补偿
                heap.add(new int[]{i,0});
            }
            for (int i = 0; i < k; i++) {
                int[] min = heap.poll();
                List<Integer> pair = new ArrayList<>(2);
                pair.add(nums1[min[0]]);
                pair.add(nums2[min[1]]);
                result.add(pair);
                if (min[1]+1 < nums2.length) { // 避免越界
                    heap.add(new int[]{min[0], min[1]+1}); // 只对 j+1
                }
//                heap.add(new int[]{min[0]+1, min[1]}); // [去重优化点] i 不动
            }
            return result;
        }

    }

    public static void main(String[] args) {
        Solution solution = new MyHeapSolution();
        solution.assertEquals(new int[]{1,7,11}, new int[]{2,4,6}, 3, new int[][]{{1,2},{1,4},{1,6}});
        solution.assertEquals(new int[]{1,1,2}, new int[]{1,2,3}, 2, new int[][]{{1,1},{1,1}});
    }
}
